import numpy as np
import cv2
def detect_fire_in_image(image_bytes: bytes) -> bytes:
    """
    输入: 图像的二进制数据
    处理: 检测火焰并标记
    输出: 处理完后的图像的二进制数据 (JPEG)
    """

    # ---------------------
    # 1. 将字节数据转换成 OpenCV 图像
    # ---------------------
    nparr = np.frombuffer(image_bytes, np.uint8)
    img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
    if img is None:
        raise ValueError("图像解码失败，请检查上传文件是否为有效的图像。")

    # ---------------------
    # 2. 转到 HSV 空间 & 设置火焰阈值 (可根据实际情况调整)
    # ---------------------
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
    lower_fire = np.array([0, 50, 50])     # [H, S, V] 下限
    upper_fire = np.array([35, 255, 255])  # [H, S, V] 上限
    mask = cv2.inRange(hsv, lower_fire, upper_fire)

    # ---------------------
    # 3. 形态学操作去噪 (这里仅作示例，参数可调)
    # ---------------------
    kernel = np.ones((5, 5), np.uint8)
    mask = cv2.erode(mask, kernel, iterations=1)
    mask = cv2.dilate(mask, kernel, iterations=2)

    # ---------------------
    # 4. 查找轮廓并标记
    # ---------------------
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    for cnt in contours:
        area = cv2.contourArea(cnt)
        if area < 100:  # 小面积过滤
            continue
        x, y, w, h = cv2.boundingRect(cnt)
        cv2.rectangle(img, (x, y), (x + w, y + h), (0, 0, 255), 2)
        cv2.putText(img, 'Fire', (x, y - 5), cv2.FONT_HERSHEY_SIMPLEX,
                    0.6, (0, 0, 255), 2)

    # ---------------------
    # 5. 将处理后的图像编码成 JPEG 格式并返回二进制数据
    # ---------------------
    _, buffer = cv2.imencode('.jpg', img)
    processed_image_bytes = buffer.tobytes()

    return processed_image_bytes


def detect_smoke_in_image(image_bytes: bytes) -> bytes:
    # ---------------------
    # 1. 将字节数据转换成 OpenCV 图像
    # ---------------------
    nparr = np.frombuffer(image_bytes, np.uint8)
    img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
    if img is None:
        raise ValueError("图像解码失败，请检查上传文件是否为有效的图像。")

    # 2. 转换为灰度图
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # 3. 阈值分割
    #   - 这里简单地设定一个固定阈值(如50)，将较暗区域视为烟雾
    #   - 实际中可能需要自适应阈值或更高级的处理
    _, thresh = cv2.threshold(gray, 50, 255, cv2.THRESH_BINARY_INV)

    # 4. 形态学操作（闭运算），去除小噪点并连接邻近区域
    kernel = np.ones((5, 5), np.uint8)
    closing = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel)

    # 5. 查找轮廓
    contours, hierarchy = cv2.findContours(
        closing, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
    )

    # 6. 遍历轮廓并过滤、绘制外接矩形
    min_area = 500  # 设定一个最小面积阈值，过滤太小的噪声
    for cnt in contours:
        area = cv2.contourArea(cnt)
        if area > min_area:
            x, y, w, h = cv2.boundingRect(cnt)
            # 在原图上绘制红色矩形框
            cv2.rectangle(img, (x, y), (x + w, y + h), (0, 0, 255), 2)

    # ---------------------
    # 7. 将处理后的图像编码成 JPEG 格式并返回二进制数据
    # ---------------------
    _, buffer = cv2.imencode('.jpg', img)
    processed_image_bytes = buffer.tobytes()

    return processed_image_bytes